package ace.cmp.spring.data.jdbc.querydsl.annotation.processor;

import com.google.auto.service.AutoService;
import com.google.common.base.CaseFormat;
import com.google.common.collect.ImmutableSet;
import com.infobip.spring.data.jdbc.annotation.processor.*;
import com.querydsl.apt.AbstractQuerydslProcessor;
import com.querydsl.apt.Configuration;
import com.querydsl.apt.ExtendedTypeFactory;
import com.querydsl.apt.TypeElementHandler;
import com.querydsl.codegen.*;
import com.querydsl.sql.codegen.NamingStrategy;
import java.lang.annotation.Annotation;
import java.util.Objects;
import java.util.Set;
import java.util.logging.Logger;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import javax.annotation.processing.Processor;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import javax.lang.model.util.ElementFilter;
import org.springframework.data.annotation.Id;
import org.springframework.data.annotation.Transient;
import org.springframework.data.relational.core.mapping.Embedded;

/**
 * @author caspar
 * @date 2023/4/3 9:22
 */
@AutoService(Processor.class)
public class LowerUnderscoreSpringDataJdbcAnnotationProcessor extends AbstractQuerydslProcessor {
  private Logger logger =
      Logger.getLogger(LowerUnderscoreSpringDataJdbcAnnotationProcessor.class.getName());
  private RoundEnvironment roundEnv;
  private CustomExtendedTypeFactory typeFactory;
  private Configuration conf;
  private final NamingStrategy namingStrategy;
  private final TypeElementHandlerFactory typeElementHandlerFactory;
  private CaseFormat projectTableCaseFormat;
  private CaseFormat projectColumnCaseFormat;

  public LowerUnderscoreSpringDataJdbcAnnotationProcessor() {
    this.namingStrategy = new AceDefaultNamingStrategy();
    this.typeElementHandlerFactory = new DefaultTypeElementHandlerFactory();
    this.projectTableCaseFormat = CaseFormat.LOWER_UNDERSCORE;
    this.projectColumnCaseFormat = CaseFormat.LOWER_UNDERSCORE;
  }

  @Override
  public Set<String> getSupportedAnnotationTypes() {
    return ImmutableSet.of(Id.class.getName());
  }

  @Override
  protected Configuration createConfiguration(RoundEnvironment roundEnv) {
    Class<? extends Annotation> entity = Id.class;
    this.roundEnv = roundEnv;
    var codegenModule = new CodegenModule();
    var typeMappings = new JavaTypeMappings();
    codegenModule.bind(TypeMappings.class, typeMappings);
    codegenModule.bind(QueryTypeFactory.class, new QueryTypeFactoryImpl("Q", "", ""));
    var springDataJdbcConfiguration =
        new SpringDataJdbcConfiguration(
            roundEnv,
            processingEnv,
            projectColumnCaseFormat,
            entity,
            null,
            null,
            Embedded.class,
            Transient.class,
            typeMappings,
            codegenModule,
            namingStrategy);
    this.conf = springDataJdbcConfiguration;
    return springDataJdbcConfiguration;
  }

  @Override
  protected TypeElementHandler createElementHandler(
      TypeMappings typeMappings, QueryTypeFactory queryTypeFactory) {
    return new TypeWithPrimaryKeyElementHandler(
        this.conf, typeFactory, typeMappings, queryTypeFactory);
  }

  @Override
  protected ExtendedTypeFactory createTypeFactory(
      Set<Class<? extends Annotation>> entityAnnotations,
      TypeMappings typeMappings,
      QueryTypeFactory queryTypeFactory) {

    var customExtendedTypeFactory =
        new AceCustomExtendedTypeFactory(
            roundEnv,
            processingEnv,
            entityAnnotations,
            typeMappings,
            queryTypeFactory,
            conf,
            processingEnv.getElementUtils(),
            projectTableCaseFormat);
    this.typeFactory = customExtendedTypeFactory;
    return customExtendedTypeFactory;
  }

  protected Set<TypeElement> collectElements() {
    var entityElements =
        roundEnv.getElementsAnnotatedWith(conf.getEntityAnnotation()).stream()
            .map(Element::getEnclosingElement)
            .filter(element -> element instanceof TypeElement)
            .map(element -> (TypeElement) element)
            .collect(Collectors.toSet());

    return entityElements.stream()
        .flatMap(this::getEntityElementWithEmbeddedEntities)
        .collect(Collectors.toSet());
  }

  private Stream<TypeElement> getEntityElementWithEmbeddedEntities(TypeElement entityElement) {
    var types = processingEnv.getTypeUtils();
    var embeddedElements =
        ElementFilter.fieldsIn(entityElement.getEnclosedElements()).stream()
            .filter(
                enclosedElement ->
                    Objects.nonNull(enclosedElement.getAnnotation(conf.getEmbeddedAnnotation())))
            .map(element -> types.asElement(element.asType()))
            .filter(element -> element instanceof TypeElement)
            .map(element -> (TypeElement) element)
            .collect(Collectors.toSet());
    return Stream.concat(Stream.of(entityElement), embeddedElements.stream());
  }
}
